#!/usr/bin/python
# coding=UTF-8
#import pysam
import argparse
from calendar import c
from math import e
from operator import ge
import sys
import os
from multiprocessing import Manager
import multiprocessing as mp
import gzip
import mmap
import pod5 as p5
import time
import datetime
import numpy as np
from ont_fast5_api.fast5_interface import get_fast5_file
from deepsignal3.utils.process_utils import fill_files_queue
from deepsignal3.utils.process_utils import get_refloc_of_methysite_in_motif
from deepsignal3.utils.process_utils import get_motif_seqs
from pathlib import Path
# from memory_profiler import profile
from deepsignal3.utils.process_utils import CIGAR2CODE
from deepsignal3.utils.process_utils import CIGAR_REGEX
from deepsignal3.utils import bam_reader
from deepsignal3.extract_features_pod5 import _group_signals_by_movetable_v2
import re
import traceback
from deepsignal3.utils.process_utils import complement_seq
from tqdm import tqdm
from typing import List, Tuple, Optional

# import threading
# 定义终止事件
# event = threading.Event()

# def get_kmer_from_header(header):
#     CO = header['CO']
#     for c in CO:
#         if c.startswith("UNC:"):
#             # 提取 UNC 后的 JSON 字符串
#             json_str = c[4:].strip()
#             # 解析 JSON 字符串
#             # 提取kmer
#             try:
#                 json_data = json.loads(json_str)
#                 model = list(json_data['models'].keys())[0]
#                 kmer = json_data['models'][model]['k']
#                 print(f"uncalled4 align model: {model}, kmer: {kmer}")

#                 return kmer
#             except KeyError as e:
#                 print(f"parse kmer by @CO:UNC err: {e}")
#                 # 存在多个UNC/多次提取的情况吗？
#                 continue
#     print(f"not found kmer from header:@CO:UNC, please check bam file first.")
#     return 0

class ExtractSignalWithKmerByUncalled4(object):
    def __init__(self, ur, ul, k):
        self.ur = self.process_ur(ur)
        self.ul = ul
        self.k = k

    # 返回成对的ur
    # 输出: [(1, 2), (3, 4), (5, 6)]
    def process_ur(self,ur):
        # ur是一串连续的开始和结束，应该长度为偶数
        assert len(ur) % 2 == 0, 'The length of ur(%d) must be an even number.' % (len(ur))
        return [(ur[i], ur[i + 1]) for i in range(0, len(ur) - 1, 2)]

    def get_clip_between_ur_ref(self, ref_range, ur_pair):
        # 返回ref开始比ur多多少，结束比ur多多少
        # 给的范围是前闭后开的
        start = ref_range[0] - ur_pair[0]
        if start > 0:
            # 是否应该直接报错呢？
            # raise Exception(f"ur({self.ur[0]}) start should not before ref start({ref_range[0]})")
            print("ur({}) should not before ref start({})".format(ur_pair[0], ref_range[0]))
            # 标记一下
            start = -1
        else:
            start = abs(start)
        end = ref_range[1] - ur_pair[1]
        if end < 0:
            # 是否应该直接报错呢？
            # raise Exception(f"ur({self.ur[1]}) end should not after ref start({ref_range[1]})")
            print("ur({}) end should not after ref start({})".format(ur_pair[1], ref_range[1]))
        return start, end

    def get_rseq_by_ur(self, rseq, ref_range):
        # 先按照ur拼接ref seq
        new_rseq_split = []
        rseq_splice = ""
        for ur_pair in self.ur:
            st,en = self.get_clip_between_ur_ref(ref_range, ur_pair)
            # st 和 en都应该大于0
            assert st >= 0 and  en >= 0, "The starting positions of ur and ref do not correspond, please check the file!"
            # 获取裁剪后的长度
            rseq_ = rseq[st: len(rseq) - en]
            assert ur_pair[1] - ur_pair[0] == len(rseq_), 'Not expected rseq(%d) and ur(%d) to have different lengths.' % (
                len(rseq_), ur_pair[1] - ur_pair[0])
            # 进行拼接
            rseq_splice = rseq_splice + rseq_
            new_rseq_split.append(rseq_)
        self.new_rseq_split = new_rseq_split
        return rseq_splice

    # 根据参考序列获取kmers
    def generate_kmers_by_rseq(self, rseq, ref_range):
        self.new_rseq = self.get_rseq_by_ur(rseq, ref_range)
        if self.k <= 0:
            raise ValueError("k must be integer greater than 0.")
        if self.k > len(self.new_rseq):
            raise ValueError("k should be smaller than length of sequence.")

        kmers = [self.new_rseq[i: i + self.k] for i in range(len(self.new_rseq) - self.k + 1)]
        return kmers

    # 获取分隔后的信号
    def get_signal_with_kmers(self, rseq, ref_range, signal):
        # 获取kmer
        kmers = self.generate_kmers_by_rseq(rseq, ref_range)
        neg_count = len([x for x in self.ul if x < 0])
        # ul减去小于0的个数
        assert len(self.ul) - neg_count == len(kmers), 'Not expected kmer len(%d) and ul(%d) to have different lengths.' % (
            len(kmers), len(self.ul) - neg_count)

        # 开始2个碱基
        first_2_nuc = kmers[0][:2]
        # 最后3个碱基
        last_3_nuc = kmers[-1][3:]
        # 最后一个kmer
        last_kmer = kmers[-1]
        # 第一个kmer
        first_kmer = kmers[0]

        # 获得信号起始位置
        signal_start = 0
        # ref 开始位置
        ref_start = self.ur[0][0]
        # 根据ur，ref本段结束的位置，记录的是ur列表中的位置
        ref_end_by_ur = 0
        # read起始位置，要不要加2呢？
        read_start = 0
        # 根据kmer和ul分隔的信号
        kmer_signal = []
        # 信号起始位置
        signal_start_pos = []
        # 碱基对应参考序列起始位置
        refs_pos = []
        # 碱基对应reads起始位置
        read_pos = []
        # kmer to signal table
        kmer_table = []

        # 把开始2个nuc的位置补上
        # 根据uncalled转tsv规则看是ur第一个值加2，因为是从第3个碱基开始(0 base)
        first_kmer_signal = []
        first_signal_start_pos = []
        for n in first_2_nuc:
            dic = {'kmer': first_kmer, 'nuc': n, 'signal': [], 'signal_pos': 0, 'ref_pos': ref_start,
                   'read_pos': read_start}
            #print(dic)
            kmer_table.append(dic)
            first_kmer_signal.append([])
            first_signal_start_pos.append(0)
            read_start += 1
            ref_start += 1

        for length in self.ul:
            # length小于0说明空了几个位置，重新定义信号开始
            if length < 0:
                signal_start += abs(length)
                continue
            segment = signal[signal_start:signal_start + length]
            signal_start_pos.append(signal_start)
            # print(f"start {signal_start} len {length}")
            # 需要注意的是，当len为0时我们这里加入的是空，但是uncalled4中是用的上一段的内容
            kmer_signal.append(list(segment))
            refs_pos.append(ref_start)
            read_pos.append(read_start)
            if ref_start + 1 == self.ur[ref_end_by_ur][1]:
                ref_end_by_ur += 1
                # 因为ref seq是拼接起来的，所以ref开始位置可能是中断的
                # read_start指定为下一段开始位置
                ref_start = self.ur[ref_end_by_ur][0]
                # dis = self.ur[ref_end_by_ur][0] - self.ur[ref_end_by_ur - 1][1] + 1
                # read_start += dis
                # read pos应该是连续的
                read_start += 1
            else:
                ref_start += 1
                read_start += 1
            # 更新开始位置
            signal_start += length

        for i in range(len(kmers)):
            dic = {'kmer':kmers[i], 'nuc':kmers[i][2], 'signal':list(kmer_signal[i]),'signal_pos':signal_start_pos[i],'ref_pos':refs_pos[i], 'read_pos':read_pos[i]}
            #print(dic)
            kmer_table.append(dic)

        # 把最后5个nuc的位置补上
        for n in last_3_nuc:
            dic = {'kmer': last_kmer, 'nuc': n, 'signal': [], 'signal_pos':signal_start_pos[-1],'ref_pos': ref_start,'read_pos': read_start}
            #print(dic)
            kmer_table.append(dic)
            kmer_signal.append([])
            signal_start_pos.append(signal_start_pos[-1])
            read_start += 1
            ref_start += 1

        kmer_signal = first_kmer_signal + kmer_signal
        signal_start_pos = first_signal_start_pos + signal_start_pos

        assert len(signal_start_pos) == len(kmers) + 5  and len(signal_start_pos) == len(kmer_signal), 'Not expected kmer len + 5(%d) , signal split by kmer len(%d) and signal start pos(%d) to have different lengths.' % (
            len(kmers) + 5  ,len(kmer_signal), len(signal_start_pos))



        assert len(self.new_rseq) == len(kmer_signal), 'Not expected new ref seq len(%d) and signal split by kmers(%d) to have different lengths.' % (
            len(self.new_rseq), len(kmer_signal))

        return kmer_table, kmer_signal

    
def extract_signal_with_kmer(read_id, chr, rseq, ref_range, signal, uncalled4):
    '''
    Args:
        read_id: read id
        chr: chr name
        rseq: 参考序列
        ref_range: 参考序列开始和结束
        signal: read对应信号
        uncalled4: uncalled4提取后的内容封装
    Returns:
        read_id
        chr
        refs_range: 新的参考基因序列起始 == ur
        tatal_signal: 字符串，signal按照kmers分隔（;）
        kmer_table: list[dic{kmer1,ref_pos,read_pos,signal}]
    '''
    kmer_table, kmer_signal = uncalled4.get_signal_with_kmers(rseq, ref_range, signal)

    # 处理每个segment signal列表
    list_segments = [",".join(map(str, segment)) for segment in kmer_signal]
    # 使用分号连接每个segment signal的字符串
    total_signal = ";".join(list_segments)

    return read_id, chr, uncalled4.ur , total_signal, kmer_table

def extract_refs_signal_with_kmer(rseq, ref_range, signal, uncalled4):
    '''
    Args:
        rseq: 参考序列
        ref_range: 参考序列开始和结束
        signal: read对应信号
        uncalled4: uncalled4提取后的内容封装
    Returns:
        ref_range: 多个ref上的起始位置，前闭后开，[(ref_start,ref_end),(ref_start,ref_end),...]
        new_rseq: 新的ref seq，将new_rseq_split拼起来的seq
        new_rseq_split: 按照ur分隔后的ref [[ref1],[ref2],...]
        kmer_signal: 按照kmers/nuc划分的信号，[[signal1],[signal2],[signal3],...]
    '''
    _, kmer_signal = uncalled4.get_signal_with_kmers(rseq, ref_range, signal)
    return uncalled4.ur, uncalled4.new_rseq, uncalled4.new_rseq_split, kmer_signal


def get_q2tloc_from_cigar(r_cigar_tuple, strand, seq_len):
    """
    insertion: -1, deletion: -2, mismatch: -3
    :param r_cigar_tuple: pysam.alignmentSegment.cigartuples
    :param strand: 1/-1 for fwd/rev
    :param seq_len: read alignment length
    :return: query pos to ref pos
    """
    fill_invalid = -2
    # get each base calls genomic position
    q_to_r_poss = np.full(seq_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar_tuple if strand == 1 else r_cigar_tuple[::-1]
    for op, op_len in cigar_ops:
        if op == 1:
            # inserted bases into ref
            for q_pos in range(curr_q_pos, curr_q_pos + op_len):
                q_to_r_poss[q_pos] = -1
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                q_to_r_poss[curr_q_pos + op_offset] = curr_r_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    q_to_r_poss[curr_q_pos] = curr_r_pos
    if q_to_r_poss[-1] == fill_invalid:
        raise ValueError(
            (
                "Invalid cigar string encountered. Reference length: {}  Cigar "
                + "implied reference length: {}"
            ).format(seq_len, curr_r_pos)
        )
    return q_to_r_poss


def _write_featurestr(write_fp, featurestr_q, control, time_wait=1):
    # print('write_process-{} starts'.format(os.getpid()))
    output = 0
    # if os.path.exists(write_fp):
    #     with open(write_fp, 'a') as wf:
    #         while True:
    #             # during test, it's ok without the sleep(time_wait)
    #             if featurestr_q.empty():
    #                 time.sleep(time_wait)
    #                 continue
    #             features_str = featurestr_q.get()
    #             if features_str == "kill":
    #                 # print('output line {}'.format(output))
    #                 # print('write_process-{} finished'.format(os.getpid()))
    #                 break
    #             for one_features_str in features_str:
    #                 output += 1
    #                 wf.write(one_features_str + "\n")
    #             wf.flush()
    # else:
    with open(write_fp, 'w') as wf:
        while True:
            # during test, it's ok without the sleep(time_wait)
            if featurestr_q.empty():
                time.sleep(time_wait)
                continue
            features_str = featurestr_q.get()
            if features_str == "kill":
                # print('output line {}'.format(output))
                # print('write_process-{} finished'.format(os.getpid()))
                break
            for one_features_str in features_str:
                output += 1
                wf.write(one_features_str + "\n")
                if control is not None and output >= control:
                    # 设定终止标志
                    # event.set()
                    sys.exit(0)
            wf.flush()

##########
# process tsv and alignment to bam and pod5
##########


def process_deepsignal(features_file, feature_Q, signal_ids, existing_keys, remove_indexes, time_wait=1, reads_per_chunk=2, qsize_limit=20):
    if features_file.endswith(".gz"):
        infile = gzip.open(features_file, 'rt')
    else:
        infile = open(features_file, 'r')
    pred_deepsignal = {}
    fea_list = []
    pre_read_id = ''
    chromosome = ''
    read_id = ''

    # encodings = ['utf-8', 'gbk', 'latin-1']
    for line in infile:
        # if features_file.endswith(".gz"):
        #    words = line.decode('gbk').strip().split("\t")
        # else:
        #    words=line.strip().split("\t")
        words = line.strip().split("\t")
        read_id = words[4]
        key = '||'.join([read_id, words[0]])
        if read_id not in signal_ids:
            continue
        if key in existing_keys or key in remove_indexes:
            continue
        if pre_read_id == '':
            pre_read_id = read_id
        elif pre_read_id != read_id:
            fea_list.append((pre_read_id, chromosome, pred_deepsignal))
            pre_read_id = read_id
            pred_deepsignal = {}
            if len(fea_list) >= reads_per_chunk:
                while feature_Q.qsize() >= qsize_limit:
                    # print('Pausing tsv input due to INPUT queue size limit. Signal_qsize=%d' %(feature_Q.qsize()), flush=True)
                    time.sleep(time_wait)
                feature_Q.put(fea_list)
                fea_list = []

        chromosome = words[0]
        pred_deepsignal[int(words[1])] = float(words[7])
    if pre_read_id != read_id:
        fea_list.append((pre_read_id, chromosome, pred_deepsignal))
    if len(fea_list) > 0:
        feature_Q.put(fea_list)
    infile.close()

# def remove_duplicates_except_first(lst):
#     seen = {}
#     result = []
#     for num in lst:
#         if num in seen:
#             result.append(None)
#             continue
#         seen[num] = True
#         result.append(num)  
#     return result
 

def align_signals(pos_pair, read_signal, ref_seq):
    # 初始化 ref_signal，长度与 ref_seq 相同
    ref_signal = [[] for _ in range(len(ref_seq))]
    
    # 记录上一个有效的 ref_pos
    last_valid_ref_pos = None

    # 遍历 pos_pair
    for i, (read_pos, ref_pos) in enumerate(pos_pair):
        if read_pos is None:
            # 如果 read_pos 为 None，表示发生删除，ref_signal 对应位置保持为空列表
            continue
        elif ref_pos is None:
            # 如果 ref_pos 为 None，表示发生插入，将信号拼接到上一个有效 ref_pos 的后面
            if last_valid_ref_pos is None:
                # 如果没有找到前一个有效的 ref_pos，插入到 ref_signal[0] 的最前面
                ref_signal[0] = read_signal[read_pos] + ref_signal[0]
            else:
                # 否则拼接到前一个有效 ref_pos 对应的 ref_signal 列表的后面
                ref_signal[last_valid_ref_pos] += read_signal[read_pos]
        else:
            # 如果 ref_pos 有效，将信号添加到 ref_signal[ref_pos]
            ref_signal[ref_pos] += read_signal[read_pos]
            # 更新 last_valid_ref_pos 为当前 ref_pos
            last_valid_ref_pos = ref_pos

    return ref_signal

def filter_n_and_update_indices(seq, pred_pos):
    # Step 1: Initialize variables
    n_count = 0
    n_positions = []
    
    # Step 2: Traverse the sequence to identify 'N' positions and count them
    for i, char in enumerate(seq):
        if char == 'N':
            n_positions.append(i)
    
    # Calculate number of 'N's to the left of each position in pred_pos
    updated_pred_pos = []
    
    for pos in pred_pos:
        # Calculate the number of 'N's before the current position
        count_n_before_pos = sum(1 for n_pos in n_positions if n_pos < pos)
        # Adjust the position by subtracting the number of 'N's before it
        updated_pos = pos - count_n_before_pos
        updated_pred_pos.append(updated_pos)
    
    # Step 3: Filter 'N' from the sequence
    filtered_seq = ''.join([char for char in seq if char != 'N'])
    return filtered_seq, updated_pred_pos, n_positions

def align_signals_and_extend_ref_seq(pos_pair, read_signal, read_seq, ref_seq,motif_seqs,methyloc,strand,ref_start,ref_end):
    # 确定最前端和最末端有效的比对索引
    first_valid_index = next((i for i, (_, ref_pos) in enumerate(pos_pair) if ref_pos is not None), len(pos_pair))
    last_valid_index = len(pos_pair) - 1 - next((i for i, (_, ref_pos) in enumerate(reversed(pos_pair)) if ref_pos is not None), len(pos_pair))

    # 生成新的 ref_seq 和 ref_signal
    new_ref_seq = []
    new_ref_signal = []

    # 填充前端未比对的部分
    for i in range(first_valid_index):
        read_pos, _ = pos_pair[i]
        if read_pos is not None:
            new_ref_seq.append(read_seq[read_pos])
            new_ref_signal.append(read_signal[read_pos])

    last_valid_ref_pos = len(new_ref_seq) - 1

    # 处理中间的比对部分
    for i in range(first_valid_index, last_valid_index + 1):
        read_pos, ref_pos = pos_pair[i]

        if ref_pos is not None:
            # 确保 new_ref_seq 的长度足够
            # while len(new_ref_seq) < ref_pos:
            #     new_ref_seq.append(None)  # 用 None 占位，表示插入的部分
            #     new_ref_signal.append([]) # 对应的信号也填充为空列表
            
            new_ref_seq.append(ref_seq[ref_pos])
            new_ref_signal.append(read_signal[read_pos] if read_pos is not None else [])
            last_valid_ref_pos = len(new_ref_seq) - 1

        elif ref_pos is None:
            if last_valid_ref_pos is not None:
                new_ref_signal[last_valid_ref_pos].extend(read_signal[read_pos])

    # 填充后端未比对的部分
    for i in range(last_valid_index + 1, len(pos_pair)):
        read_pos, _ = pos_pair[i]
        if read_pos is not None:
            new_ref_seq.append(read_seq[read_pos])
            new_ref_signal.append(read_signal[read_pos])
            

    new_ref_seq = ''.join(base  for base in new_ref_seq)
    ref_readlocs = dict()
    ref_poss = []
    pred_pos = []
    ref_pos = -1
    tsite_locs = get_refloc_of_methysite_in_motif(
        new_ref_seq, set(motif_seqs), methyloc)
    for loc_in_read in tsite_locs:
        if loc_in_read<first_valid_index:
            ref_pos = -1
            ref_poss.append(ref_pos)
            pred_pos.append(loc_in_read)
            continue
        if loc_in_read>last_valid_index:
            ref_pos = -1
            ref_poss.append(ref_pos)
            pred_pos.append(loc_in_read)
            continue
        if strand == "-":
            ref_pos = ref_end-loc_in_read-1+first_valid_index
        else:
            ref_pos = ref_start+loc_in_read-first_valid_index
        #ref_readlocs[loc_in_read+first_valid_index] = ref_pos
        ref_poss.append(ref_pos)
        pred_pos.append(loc_in_read)
    #new_ref_seq, pred_pos,n_positions=filter_n_and_update_indices(new_ref_seq, pred_pos)
    ref_readlocs = dict(zip(pred_pos, ref_poss))
    #new_ref_signal = [new_ref_signal[i] for i in range(len(new_ref_signal)) if i not in n_positions]
    #n_lens=len(n_positions)
    return new_ref_seq, new_ref_signal, ref_readlocs, ref_poss, pred_pos#, n_lens
# def get_reference_length_from_cigar(cigar_tuples):
#     ref_length = 0
#     insertions = 0
#     for op, length in cigar_tuples:
#         if op in [0, 2, 3, 7, 8]:  # M, D, N, =, X
#             ref_length += length
#         elif op == 1:  # I
#             insertions += length
#     return ref_length, insertions

# def get_reference_length_from_md(md_tag):
#     ref_length = 0
#     for match in re.findall(r'(\d+)|([A-Z])|(\^[A-Z]+)', md_tag):
#         if match[0]:
#             ref_length += int(match[0])
#         elif match[1]:
#             ref_length += 1
#         elif match[2]:
#             ref_length += len(match[2][1:])  # 删除序列长度
#     return ref_length
def get_md_reference_length(md_tag):
    """
    计算 MD 标签描述的参考序列长度。
    
    参数:
    - md_tag: str - MD 标签字符串（如 '48T42G8' 或 '5^A5'）
    
    返回:
    - int - 参考序列总长度
    """
    total_length = 0
    # 提取所有数字（匹配长度）
    matches = re.findall(r'\d+', md_tag)
    total_length += sum(int(m) for m in matches)
    # 提取所有单字符不匹配（每个贡献长度 1）
    mismatches = re.findall(r'([A-Z])', md_tag)
    total_length += len(mismatches)
    # 提取所有删除序列（每个删除序列的长度贡献到总长度）
    deletions = re.findall(r'\^(\w+)', md_tag)
    total_length += sum(len(d) for d in deletions)
    return total_length

# def build_alignment_sequence(alignment):
#     """
#     根据 CIGAR 和查询序列构建扩展序列。
    
#     参数:
#     - alignment: pysam.AlignedSegment - BAM 对齐记录
    
#     返回:
#     - str - 扩展序列（包含匹配、不匹配、插入和删除）
#     """
#     if not alignment.has_tag('MD'):
#         return None
#     pairs = alignment.get_aligned_pairs(with_seq=True)
#     s = []
#     for qpos, rpos, qbase in pairs:
#         if qpos is not None and rpos is not None:
#             # 匹配或不匹配：添加查询序列的碱基
#             s.append(qbase)
#         elif qpos is None and rpos is not None:
#             # 删除：添加 '-'
#             s.append('-')
#         elif qpos is not None and rpos is None:
#             # 插入：添加小写查询序列的碱基
#             s.append(qbase.lower())
#     return ''.join(s)



# def get_alignment_length(alignment):
#     # 定义常量以提高可读性
#     BAM_CSOFT_CLIP = 4  # 软剪切操作代码
#     BAM_CHARD_CLIP = 5  # 硬剪切操作代码
#     """
#     计算对齐长度，排除软剪切和硬剪切操作。
    
#     参数:
#     alignment (pysam.AlignedSegment): BAM 文件的对齐记录。
    
#     返回:
#     int: 总对齐长度，排除软剪切和硬剪切。
#     """
#     if alignment is None or alignment.cigartuples is None:
#         return 0
#     total_length = 0
#     for op, length in alignment.cigartuples:
#         if op not in (BAM_CSOFT_CLIP, BAM_CHARD_CLIP):
#             total_length += length
#     return total_length

# def check_md_validity(alignment):
#     """
#     检查 MD 标签的有效性。
    
#     参数:
#     - alignment: pysam.AlignedSegment - BAM 对齐记录
    
#     抛出:
#     - AssertionError - 如果 MD 标签无效
#     """
#     # if not alignment.has_tag('MD'):
#     #     raise AssertionError("MD tag not found in read {}.".format(alignment.query_name))
#     md_tag = alignment.get_tag('MD')
#     md_len = get_md_reference_length(md_tag)
#     max_len = get_alignment_length(alignment)#sum(length for op, length in alignment.cigartuples if op in [0, 1, 2, 6, 7, 8])
#     # 计算插入数（小写字符的数量）
#     #insertions = sum(1 for c in s if c.islower())
#     insertions = sum(length for op, length in alignment.cigartuples if op in [1, 6])
#     if md_len + insertions > max_len:
#         raise AssertionError(
#             f"Invalid MD tag: MD length {md_len} mismatch with CIGAR length {max_len} and {insertions} insertions"
#         )
#     return
def get_md_reference_length(md_tag: str) -> int:
    """Calculate the reference sequence length from an MD tag."""
    length = 0
    md_idx = 0
    nmatches = 0
    
    while md_idx < len(md_tag):
        if md_tag[md_idx].isdigit():
            nmatches *= 10
            nmatches += int(md_tag[md_idx])
            md_idx += 1
            continue
        else:
            length += nmatches
            nmatches = 0
            if md_tag[md_idx] == '^':
                md_idx += 1
                while md_idx < len(md_tag) and md_tag[md_idx].isupper():
                    md_idx += 1
                    length += 1
            else:
                md_idx += 1
                length += 1
    
    length += nmatches
    return length

def get_alignment_length(cigartuples: List[Tuple[int, int]]) -> int:
    """Calculate the reference sequence length from CIGAR tuples."""
    length = 0
    for op, op_len in cigartuples:
        if op in {0, 2, 3, 7, 8}:  # CMATCH, CDEL, CREF_SKIP, CEQUAL, CDIFF
            length += op_len
    return length

def count_insertions(alignment_sequence: str) -> int:
    """Count the number of insertions (lowercase letters) in the alignment sequence."""
    return sum(1 for c in alignment_sequence if c.islower())

def check_md_cigar_length(md_tag: str, cigartuples: List[Tuple[int, int]], alignment_sequence: str) -> None:
    """
    Check if MD tag length + insertions exceeds CIGAR reference length.
    
    Args:
        md_tag: The MD tag string (e.g., "5^C5").
        cigartuples: List of CIGAR tuples [(op, length), ...].
        alignment_sequence: The alignment sequence containing insertions as lowercase letters.
    
    Raises:
        AssertionError: If MD length + insertions > CIGAR reference length.
    """
    md_len = get_md_reference_length(md_tag)
    cigar_ref_len = get_alignment_length(cigartuples)
    insertions = count_insertions(alignment_sequence)
    
    if md_len + insertions > cigar_ref_len:
        raise AssertionError(
            f"Invalid MD tag: MD length {md_len} + {insertions} insertions "
            f"exceeds CIGAR length {cigar_ref_len}"
        )
def process_dorado(bam_index, pod5_dr, bisulfite, read_ids_Q, output_Q, motif_seqs,progress_bar,lock,process_chr,use_ref,label, methyloc=0, time_wait=1,mapq=0,kmer=6, d_batch_size=2, qsize_limit=22):
    fea_list = []
    while True:
        while read_ids_Q.empty():
            time.sleep(time_wait)
        read_data = read_ids_Q.get()
        if read_data == "kill":
            read_ids_Q.put("kill")
            break
        for read_name in read_data:
            while output_Q.qsize() >= qsize_limit:
                # print('Pausing bam and pod5 input due to OUTPUT queue size limit. Output_qsize=%d' %(output_Q.qsize()), flush=True)
                time.sleep(time_wait)
            with lock:
                progress_bar.update(1)
            read = pod5_dr.get_read(read_name)
            if read is None:
                continue
            signal = read.signal
            shift_dacs_to_pa = read.calibration.offset
            scale_dacs_to_pa = read.calibration.scale
            # read_iter=bam_index.find(read_name)
            try:
                for bam_read in bam_index.get_alignments(read_name):
                    if bam_read.is_supplementary or bam_read.is_secondary:
                        continue
                    reference_name = bam_read.reference_name
                    if reference_name is None:
                        continue
                    # if reference_name is not None:
                    #     reference_name=reference_name.lower()
                    if bam_read.mapping_quality < mapq:
                        continue

                    if process_chr is not None:
                        if isinstance(process_chr, list):
                            if reference_name not in process_chr:
                                continue
                        else:
                            if process_chr[:2]!='no':
                                if reference_name != process_chr:
                                    continue
                            elif process_chr[:2]=='no':
                                if reference_name == process_chr[2:]:
                                    continue

                    seq = bam_read.get_forward_sequence().upper()
                    if seq is None or signal is None:
                        continue
                    
                    if use_ref:
                        if not bam_read.has_tag('MD'):
                            print('not have MD of readid {}'.format(read_name))
                            continue
                        #md_tag = bam_read.get_tag('MD')
                        # cigar_ref_length, insertions = get_reference_length_from_cigar(bam_read.cigartuples)
                        # md_ref_length = get_reference_length_from_md(md_tag)
                        # if md_ref_length + insertions < cigar_ref_length:
        
                        #     print('error md_ref_length {} cigar_ref_length {} insertions {} of readid {}'.format(md_ref_length, cigar_ref_length, insertions, read_name))
                        #     continue
                        # try:
                        #     #check_md_validity(bam_read)
                        #     md_tag = bam_read.get_tag("MD")  # 获取 MD 标签
                        #     cigartuples = bam_read.cigartuples  # 获取 CIGAR tuples
                        #     alignment_sequence = bam_read.query_alignment_sequence
                        #     check_md_cigar_length(md_tag, cigartuples, alignment_sequence)
                        #     print(f"Read {bam_read.query_name} passed MD/CIGAR check.", flush=True)
                        # except AssertionError as e:
                        #     #print(f"Error in {bam_read.query_name}: {e}")
                        #     #print(f"Error in {read_name}: {e}")
                        #     print(f"Read {bam_read.query_name} failed: {e}", flush=True)
                        #     continue
                        try:
                            if bam_read.is_reverse:
                                ref_seq = complement_seq(bam_read.get_reference_sequence().upper())
                            else:
                                ref_seq = bam_read.get_reference_sequence().upper()
                        except Exception as e:
                            print(f"Read {bam_read.query_name} failed: {e}", flush=True)
                            continue                   
                    
                    if bam_read.infer_query_length() != len(seq):
                        print(
                            'infer read length is not same as seq length of readid {}'.format(read_name))
                        continue
                    read_dict = dict(bam_read.tags)

                    strand_code = 0 if bam_read.is_reverse else 1

                    strand = "-" if strand_code == 0 else "+"
                    # find_key=(read_name,reference_name)
                    ref_start = bam_read.reference_start
                    ref_end = bam_read.reference_end
                    cigar_tuples = bam_read.cigartuples

                    num_trimmed = read_dict["ts"]
                    if bam_read.has_tag('sp'):
                        num_trimmed += bam_read.get_tag('sp')
                    shift_pa_to_norm = read_dict["sm"]
                    scale_pa_to_norm = read_dict["sd"]
                    mv_table = read_dict["mv"]
                    if num_trimmed >= 0:
                        # (signal[num_trimmed:] - norm_shift) / norm_scale
                        signal_trimmed = signal[num_trimmed:]
                    else:
                        # (signal[:num_trimmed] - norm_shift) / norm_scale
                        signal_trimmed = signal[:num_trimmed]
                    # sshift, sscale = np.mean(signal_trimmed), float(np.std(signal_trimmed))
                    # if sscale == 0.0:
                    #    norm_signals = signal_trimmed
                    # else:
                    #    norm_signals = (signal_trimmed - sshift) / sscale
                    signal_group = _group_signals_by_movetable_v2(
                        signal_trimmed, np.asarray(mv_table[1:]), int(mv_table[0]))
                    if use_ref:
                        if bam_read.has_tag('ur') and bam_read.has_tag('ul'):
                            # ur是参考序列中的起始位置，中断暂时没遇到
                            ur = bam_read.get_tag("ur")
                            # ul是信号对应kmer
                            ul = bam_read.get_tag("ul")
                            ref_range = [ref_start, ref_end]
                            uncalled4_obj = ExtractSignalWithKmerByUncalled4(ur, ul, kmer)
                            ref_range, seq,rseqs, signal_group = extract_refs_signal_with_kmer(ref_seq, ref_range, signal, uncalled4_obj)
                            #ref_readlocs=dict()
                            ref_poss=[]
                            pred_pos=[]
                            ref_len=[]
                            for i in range(len(rseqs)):
                                rseq=rseqs[i]
                                ref_start, ref_end = ref_range[i]
                                ref_len.append(len(rseq))
                                tsite_locs = get_refloc_of_methysite_in_motif(
                                    rseq, set(motif_seqs), methyloc)
                                for loc_in_read in tsite_locs:
                                    if strand == "-":
                                        ref_pos = ref_end - loc_in_read - 1
                                    else:
                                        ref_pos = ref_start + loc_in_read
                                    ref_poss.append(ref_pos)
                                    # if i!=0:
                                    #     pred_pos.append(loc_in_read+sum(ref_len[:i]))
                                    # else:
                                    #     pred_pos.append(loc_in_read)
                                    pred_pos.append(loc_in_read+sum(ref_len[:i]))
                                ref_readlocs = dict(zip(pred_pos, ref_poss))

                        else:
                            pos_pair=[]
                            for read_pos,ref_pos in bam_read.get_aligned_pairs():
                                # if ref_pos is None:
                                #     continue
                                if read_pos is None:
                                    if bam_read.is_reverse:
                                        pos_pair.append((None,ref_end-ref_pos-1))
                                    else:
                                        pos_pair.append((None,ref_pos-ref_start))
                                    continue
                                if ref_pos is None:
                                    if bam_read.is_reverse:
                                        pos_pair.append((len(seq)-read_pos-1,None))
                                    else:
                                        pos_pair.append((read_pos,None))
                                    continue
                                if bam_read.is_reverse:
                                    pos_pair.append((len(seq)-read_pos-1,ref_end-ref_pos-1))
                                else:
                                    pos_pair.append((read_pos,ref_pos-ref_start))
                            #pos_pair=bam_read.get_aligned_pairs()
                            #signal_group=align_signals(pos_pair,signal_group,ref_seq)
                            if strand == "-":
                                pos_pair.reverse()
                            seq,signal_group, ref_readlocs, ref_poss, pred_pos=align_signals_and_extend_ref_seq(pos_pair, signal_group, seq, ref_seq,motif_seqs,methyloc,strand,ref_start,ref_end)
                    # if use_ref:
                    #     seq=ref_seq
                        # print(
                        #     '{},{}'.format(n_lens,read_name))
                    if len(signal_group) != len(seq):
                        print('signal to seq error!')
                        continue
                    if use_ref is False:
                        qalign_start = bam_read.query_alignment_start
                        qalign_end = bam_read.query_alignment_end
                        if bam_read.is_reverse:
                            seq_start = len(seq) - qalign_end
                            seq_end = len(seq) - qalign_start
                        else:
                            seq_start = qalign_start
                            seq_end = qalign_end
                        q_to_r_poss = get_q2tloc_from_cigar(
                            cigar_tuples, strand_code, (seq_end - seq_start)
                        )
                        ref_readlocs = dict()
                        ref_poss = []
                        pred_pos = []
                        tsite_locs = get_refloc_of_methysite_in_motif(
                            seq, set(motif_seqs), methyloc)
                        for loc_in_read in tsite_locs:
                            if seq_start <= loc_in_read < seq_end:
                                offset_idx = loc_in_read - seq_start
                                if q_to_r_poss[offset_idx] != -1:
                                    if strand == "-":
                                        # pos = '.'#loc_in_read
                                        ref_pos = ref_end - 1 - \
                                            q_to_r_poss[offset_idx]
                                    else:
                                        # pos = loc_in_read
                                        ref_pos = ref_start + \
                                            q_to_r_poss[offset_idx]
                                    ref_readlocs[loc_in_read] = ref_pos
                                    ref_poss.append(ref_pos)
                                    pred_pos.append(loc_in_read)
                                else:
                                    ref_pos = -1
                                    ref_readlocs[loc_in_read] = ref_pos
                                    ref_poss.append(ref_pos)
                                    pred_pos.append(loc_in_read)
                            else:
                                ref_pos = -1
                                ref_readlocs[loc_in_read] = ref_pos
                                ref_poss.append(ref_pos)
                                pred_pos.append(loc_in_read)
                       
                    if len(ref_poss) == 0:
                        continue
                    bisulfite_ref = dict()
                    # pred_pos=dict()
                    if bisulfite is not None:
                        flag=0                       
                        for pos in pred_pos:
                            ref_pos = ref_readlocs[pos]
                            key_combine = '||'.join([reference_name, str(ref_pos)])
                            if key_combine in bisulfite:
                                bisulfite_ref[pos] = bisulfite[key_combine]
                                flag=1
                            else:
                                bisulfite_ref[pos] = -1
                        if flag==0:
                            continue
                    # for ref_pos in ref_readlocs.keys():
                    #     key_combine='||'.join([reference_name,str(ref_pos)])
                    #     if key_combine in bisulfite:
                    #         if seq[ref_readlocs[ref_pos]]=='C':
                    #             bisulfite_ref[ref_pos]=bisulfite[key_combine]
                    #             #pred_pos[ref_pos]=ref_readlocs[ref_pos]
                    #             ref_poss.append(ref_pos)
                    
                    
                    signal_group_new = []
                    for sig in signal_group:
                        signal_group_new.append(
                            np.round(np.array(sig), decimals=6))

                    norm_signals_text = ';'.join(
                        [",".join([str(y) for y in x]) for x in signal_group_new])
                    pred_deepsignal_text = '.'

                    pred_dorado_text = '.'
                    mean_pred_text = '.'
                    pred_label_text = '.'
                    pred_pos_text = ','.join(
                        [str(x) for x in pred_pos])
                    sample_id = '\t'.join(
                        [bam_read.query_name, str(reference_name), str(ref_start)])
                    if len(bisulfite_ref)==0 and label is None:
                        bisulfite_text = '.'
                    else:
                        if len(bisulfite_ref)!=0:
                            bisulfite_text = ','.join(
                            [str(bisulfite_ref[x]) for x in pred_pos])
                        else:
                            bisulfite_text = ','.join(
                            [str(label*100) for x in pred_pos])
                    ref_pos_text = ','.join(
                        [str(ref_readlocs[x]) for x in pred_pos])
                    
                    fea_str = '\t'.join([sample_id, seq, norm_signals_text, pred_pos_text, pred_dorado_text, pred_deepsignal_text, mean_pred_text, pred_label_text, str(bam_read.mapping_quality),
                                         str(shift_dacs_to_pa), str(scale_dacs_to_pa), str(shift_pa_to_norm), str(scale_pa_to_norm), bisulfite_text, ref_pos_text])
                    fea_list.append(fea_str)
                    if len(fea_list) >= d_batch_size:
                        output_Q.put(fea_list)
                        fea_list = []
            except Exception as e:
                traceback.print_exc()
                continue
    if len(fea_list) > 0:
        output_Q.put(fea_list)


def generate_key(line):
    # 根据前两列生成key
    columns = line.split()[:2]
    return '||'.join(columns)


def remove_last_line(file_path):
    # 打开原始文件以读取模式
    with open(file_path, 'r') as f:
        # 定位到文件末尾
        f.seek(0, os.SEEK_END)

        # 获取文件大小
        file_size = f.tell()

        # 如果文件为空，直接返回
        if file_size == 0:
            return

        # 逐行向前查找，直到找到倒数第二个换行符位置
        pos = file_size - 1
        while pos > 0 and f.read(1) != '\n':
            pos -= 1

        # 截断文件到倒数第二个换行符位置
        f.seek(pos, os.SEEK_SET)
        f.truncate()


def process_file(output_file):
    # 检查输出文件是否存在
    try:
        with open(output_file, 'r') as f:
            existing_keys = set()
            for line in f:
                existing_keys.add(generate_key(line))
    except FileNotFoundError:
        existing_keys = set()
    return existing_keys


def remove_key(key_input):
    remove_indexes = set()
    with open(key_input, 'r') as input_file:
        for line in input_file:
            key = line.strip()
            remove_indexes.add(key)
    return remove_indexes


def read_position_file(position_file):
    key_sep = "||"
    postions = set()
    with open(position_file, "r") as rf:
        for line in rf:
            words = line.strip().split()
            postions.add(key_sep.join(words[:2]))
    return postions


# def read_bed(bisulfite_bed):
#     key_sep = "||"
#     freqinfo = {}
#     with open(bisulfite_bed, "r") as rf:
#         for line in rf:
#             words = line.strip().split()
#             if int(words[9]) < 5:
#                 continue
#             m_key = key_sep.join([words[0], words[1]])
#             freqinfo[m_key] = float(words[10])
#     return freqinfo

def read_bed(bisulfite_bed, strict=False,depth_threshold=5,label=1):
    key_sep = "||"
    depth_info = {}
    value_info = {}

    with open(bisulfite_bed, "r") as rf:
        for line in rf:
            words = line.strip().split()
            chrom = words[0]#.lower()
            pos = int(words[1])
            if len(words)==6:
                strand=words[5]
                #depth = int(words[4])
                value = label*100.0
            else:
                strand = words[5]
                depth = int(words[9])
                value = float(words[10])

                if depth < depth_threshold:
                    continue

            # 使用染色体、位置和链方向来生成唯一的key
            m_key = key_sep.join([chrom, str(pos), strand])
            #depth_info[m_key] = depth
            value_info[m_key] = value

    freqinfo = {}

    if strict:
        # 严格模式下，检查正负链是否匹配，并筛选满足新条件的点
        for m_key in value_info:
            chrom, pos, strand = m_key.split(key_sep)
            pos = int(pos)
            if strand == '+':
                pair_key = key_sep.join([chrom, str(pos + 1), '-'])
            else:
                pair_key = key_sep.join([chrom, str(pos - 1), '+'])
            
            # 检查匹配的key是否存在
            if pair_key in value_info:
                # 取得正负链的值
                value_1 = value_info[m_key]
                value_2 = value_info[pair_key]
                
                # 筛选条件：正负链要么都大于等于95，要么都小于等于5
                if value_1 >= 95 and value_2 >= 95:
                    # 将key转换为不包含正负链的形式
                    freq_key = key_sep.join([chrom, str(pos)])
                    freqinfo[freq_key] = 100
                elif value_1 <= 5 and value_2 <= 5:
                    freq_key = key_sep.join([chrom, str(pos)])
                    freqinfo[freq_key] = 0
    else:
        # 非严格模式下，直接记录所有满足深度要求的点
        for m_key in value_info:
            chrom, pos, strand = m_key.split(key_sep)
            pos = int(pos)
            freq_key = key_sep.join([chrom, str(pos)])
            freqinfo[freq_key] = value_info[m_key]

    return freqinfo




def read_tsv(key_input):
    key_indexes = []
    with open(key_input, 'r') as input_file:
        for line in input_file:
            key = line.strip()
            key_indexes.append(key)
    return key_indexes

def read_id(pod5_dir):
    # pod5_dr=pod5.DatasetReader(pod5_dir, recursive=True)
    # ids_set=set(pod5_dr.read_ids)
    ids_set = list()
    recursive = True
    glob = Path.rglob if recursive else Path.glob
    # with pod5.DatasetReader(pod5_dir, recursive=True,max_cached_readers=1) as dataset:
    #     for read_record in tqdm(dataset,desc="read ids"):
    #         ids_set.add(read_record.read_id)
    for pod5 in glob(Path(pod5_dir), "*.pod5"):
        with p5.DatasetReader(pod5, recursive=True, max_cached_readers=1) as dataset:
            for read_record in dataset:
                ids_set.append(str(read_record.read_id))

    print(len(ids_set))
    return ids_set

def process_chr(chr_file):
    chr_indexes = []
    with open(chr_file, 'r') as input_file:
        for line in input_file:
            chr = line.strip().split(',')
            for i in chr:
                chr_indexes.append(i)
    return chr_indexes

def extract(args):
    manager = mp.Manager()
    output_Q = manager.Queue()
    # bam_file=pysam.AlignmentFile(args.bam,'rb',check_sq=False,ignore_truncation=True)
    # print('%s: Building BAM index.' %str(datetime.datetime.now()), flush=True)
    # pysam.IndexedReads(bam_file,multiple_iterators=True)
    bam_index = bam_reader.ReadIndexedBam(args.bam)
    # bam_index.build()
    # print('%s: Finished building BAM index.' %str(datetime.datetime.now()), flush=True)
    signal_dir = os.path.abspath(args.signal)
    if os.path.isdir(signal_dir):
        pod5_dr = p5.DatasetReader(signal_dir, recursive=True,threads=args.pod5_proc)#, index=True #May cause the program to freeze
    else:
        pod5_dr = p5.DatasetReader(signal_dir, threads=args.pod5_proc)#, index=True #May cause the program to freeze
    if args.tsv is not None:
        signal_ids = read_tsv(args.tsv)
    else:
        if os.path.isdir(signal_dir):
            signal_ids = read_id(args.signal)
        else:
            signal_ids = list(pod5_dr.read_ids)
    read_ids_Q = manager.Queue()
    fill_files_queue(read_ids_Q, signal_ids, 20, True)
    #read_ids_Q.put("kill")

    # existing_keys=process_file(args.write_path)

    bisulfite = None
    if args.bed is not None:
        bisulfite = read_bed(args.bed,args.strict,args.depth,args.label)
        print('read bisulfite')
    timewait = args.timewait
    ex_dp = args.nproc - 1
    ex_procs = []
    is_dna = False if args.rna else True
    motif_seqs = get_motif_seqs(args.motifs, is_dna)
    print('get motifs')
    lock = manager.Lock()
    
    if args.chr is not None:
        if os.path.isfile(args.chr):
            chr_indexes = process_chr(args.chr)
        else:
            chr_indexes = args.chr.split(',')
    else:
        chr_indexes = None
    with tqdm(total=len(signal_ids), desc="Processing", ncols=100) as progress_bar:
        for i in range(ex_dp):
            pb = mp.Process(target=process_dorado, args=(bam_index, pod5_dr, bisulfite, read_ids_Q, output_Q, motif_seqs,progress_bar,lock,chr_indexes,args.ref,args.label, args.mod_loc, timewait,args.mapq,args.kmer),
                            name="pb_reader")

            pb.daemon = True
            pb.start()
            ex_procs.append(pb)
    print('start process dorado')
    p_w = mp.Process(target=_write_featurestr, args=(args.write_path, output_Q, args.control, timewait),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    print('start writer')
    # while args.control and not event.is_set():
    #    sys.exit(0)
    read_ids_Q.put("kill")
    for pb in ex_procs:
        pb.join()
    output_Q.put("kill")
    print('finished')
    p_w.join()


def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--keyset", type=str, required=False)
    parser.add_argument("--hp", type=str, required=False,
                        help='high confidence positive position')
    parser.add_argument("--hn", type=str, required=False,
                        help='high confidence negative position')
    parser.add_argument("--bam", type=str, required=True)
    parser.add_argument("--pod5", action="store_true", default=False, required=False,
                        help='use pod5, default false')
    parser.add_argument("--bed", type=str, required=False)
    parser.add_argument("--signal", type=str, required=True)
    parser.add_argument("--tsv", type=str, required=False)
    parser.add_argument("--write_path", type=str, required=True)
    parser.add_argument("--nproc", "-p", type=int, required=True)
    parser.add_argument("--timewait", "-t", default=0.1,
                        type=float, required=False)
    parser.add_argument("--d_batch_size", action="store", type=int, default=2,
                        required=False)
    parser.add_argument("--control", type=int, required=False,
                        help='test')
    parser.add_argument("--label",  type=int,
                        required=False)
    parser.add_argument('--mapq', type=int, default=0, required=False)
    parser.add_argument('--depth', type=int, default=5, required=False)
    parser.add_argument('--kmer', type=int, default=6, required=False)
    parser.add_argument('--pod5_proc','-d', type=int, default=10, required=False)
    parser.add_argument(
        "--motifs",
        action="store",
        type=str,
        required=False,
        default="CG",
        help="motif seq to be extracted, default: CG. "
        "can be multi motifs splited by comma "
        "(no space allowed in the input str), "
        "or use IUPAC alphabet, "
        "the mod_loc of all motifs must be "
        "the same",
    )
    parser.add_argument(
        "--mod_loc",
        action="store",
        type=int,
        required=False,
        default=0,
        help="0-based location of the targeted base in the motif, default 0",
    )
    parser.add_argument(
        "--rna",
        action="store_true",
        default=False,
        required=False,
        help="the fast5 files are from RNA samples. if is rna, the signals are reversed. "
        "NOTE: Currently no use, waiting for further extentsion",
    )
    parser.add_argument(
        "--ref",
        action="store_true",
        default=False,
        required=False,
        help="wetheter use ref sequence  instead of read sequence, default false.",
    )
    parser.add_argument('--strict', action="store_true", default=False, required=False)
    parser.add_argument('--test', action="store_true", default=False, required=False)
    parser.add_argument("--chr", type=str, required=False,help='only extract chr')

    return parser.parse_args()


def main():
    args = parse_args()
    extract(args)


if __name__ == '__main__':
    sys.exit(main())
